#include <cstdio>
#include <cstdlib>
#include <png.h>
#include <cuda_runtime.h>
#include <math.h>
#include <algorithm>

#define WIDTH 1024*4
#define HEIGHT 1024*4
#define MAX_DWELL 512
#define RUNS 10

struct Complex {
    float re, im;
    __device__ Complex(float r, float i) : re(r), im(i) {}
    __device__ float magnitude2() const { return re*re + im*im; }
    __device__ Complex operator*(const Complex &b) const {
        return Complex(re*b.re - im*b.im, im*b.re + re*b.im);
    }
    __device__ Complex operator+(const Complex &b) const {
        return Complex(re + b.re, im + b.im);
    }
};

// CUDA kernel: 计算每个像素的dwell（逃逸步数）
__global__ void mandelbrot_kernel(int *dwells, int width, int height,
                                  float re_min, float re_max,
                                  float im_min, float im_max) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;
    if (x >= width || y >= height) return;

    // 计算像素对应复平面坐标
    float re = re_min + (re_max - re_min) * x / (width - 1);
    float im = im_min + (im_max - im_min) * y / (height - 1);

    Complex c(re, im);
    Complex z(0.f, 0.f);

    int dwell = 0;
    while (dwell < MAX_DWELL && z.magnitude2() < 4.0f) {
        z = z * z + c;
        dwell++;
    }

    dwells[y * width + x] = dwell;
}

// 用dwell生成颜色：黑->蓝->白渐变
void dwell_to_color(int dwell, unsigned char &r, unsigned char &g, unsigned char &b) {
    if (dwell >= MAX_DWELL) {
        r = g = b = 0; // Mandelbrot集内部黑色
    } else {
        int cutoff = MAX_DWELL / 4;
        if (dwell < 0) dwell = 0;
        if (dwell <= cutoff) {
            r = 0;
            g = 0;
            b = 128 + dwell * 127 / cutoff;
        } else {
            b = 255;
            int val = (dwell - cutoff) * 255 / (MAX_DWELL - cutoff);
            r = g = val;
        }
    }
}

// 保存PNG图片
void save_png(const char* filename, int *dwells, int width, int height) {
    FILE *fp = fopen(filename, "wb");
    if (!fp) {
        fprintf(stderr, "Failed to open file for writing\n");
        exit(1);
    }
    png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
    if (!png_ptr) {
        fprintf(stderr, "png_create_write_struct failed\n");
        exit(1);
    }
    png_infop info_ptr = png_create_info_struct(png_ptr);
    if (!info_ptr) {
        fprintf(stderr, "png_create_info_struct failed\n");
        exit(1);
    }
    if (setjmp(png_jmpbuf(png_ptr))) {
        fprintf(stderr, "Error during PNG creation\n");
        exit(1);
    }
    png_init_io(png_ptr, fp);

    // 写入头信息
    png_set_IHDR(png_ptr, info_ptr, width, height,
                 8, PNG_COLOR_TYPE_RGB, PNG_INTERLACE_NONE,
                 PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
    png_write_info(png_ptr, info_ptr);

    // 写入像素数据
    png_bytep row = (png_bytep) malloc(3 * width * sizeof(png_byte));
    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            unsigned char r, g, b;
            dwell_to_color(dwells[y * width + x], r, g, b);
            row[x * 3 + 0] = r;
            row[x * 3 + 1] = g;
            row[x * 3 + 2] = b;
        }
        png_write_row(png_ptr, row);
    }
    free(row);

    png_write_end(png_ptr, nullptr);
    png_destroy_write_struct(&png_ptr, &info_ptr);
    fclose(fp);
}

int main() {
    int *d_dwells;
    size_t size = WIDTH * HEIGHT * sizeof(int);
    int *h_dwells = (int*)malloc(size);
    if (!h_dwells) {
        fprintf(stderr, "Failed to allocate host memory\n");
        return 1;
    }

    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);

    float total_time = 0.0f;
    size_t global_peak_mem = 0;

    dim3 block(16, 16);
    dim3 grid((WIDTH + block.x -1) / block.x, (HEIGHT + block.y -1) / block.y);

    float re_min = -2.0f, re_max = 1.0f;
    float im_min = -1.5f, im_max = 1.5f;

    for (int i = 0; i < RUNS; ++i) {
        // 每次循环都重新申请和释放显存，保证测试一致性
        size_t free_mem_before, free_mem_after;
        cudaMemGetInfo(&free_mem_before, nullptr);

        cudaMalloc(&d_dwells, size);

        cudaMemGetInfo(&free_mem_after, nullptr);
        size_t allocated_mem = free_mem_before - free_mem_after;

        size_t peak_mem = allocated_mem; // 初始化峰值为分配后占用

        // kernel调用前记录显存峰值
        size_t mem_before_kernel;
        cudaMemGetInfo(&mem_before_kernel, nullptr);
        peak_mem = std::max(peak_mem, free_mem_before - mem_before_kernel);

        cudaEventRecord(start);
        mandelbrot_kernel<<<grid, block>>>(d_dwells, WIDTH, HEIGHT, re_min, re_max, im_min, im_max);
        cudaEventRecord(stop);
        cudaEventSynchronize(stop);

        // kernel调用后显存
        size_t mem_after_kernel;
        cudaMemGetInfo(&mem_after_kernel, nullptr);
        peak_mem = std::max(peak_mem, free_mem_before - mem_after_kernel);

        float ms;
        cudaEventElapsedTime(&ms, start, stop);

        // 复制回CPU
        cudaMemcpy(h_dwells, d_dwells, size, cudaMemcpyDeviceToHost);

        // 释放显存
        cudaFree(d_dwells);

        // 释放后检测显存峰值
        size_t free_mem_end;
        cudaMemGetInfo(&free_mem_end, nullptr);
        peak_mem = std::max(peak_mem, free_mem_before - free_mem_end);
        global_peak_mem = std::max(global_peak_mem, peak_mem);

        printf("Run %3d: GPU time = %.6f ms, Peak memory usage = %.2f MB\n", i, ms, peak_mem / (1024.0 * 1024.0));

        // 统计第1~99次时间
        if (i > 0) total_time += ms;
    }

    save_png("mandelbrot.png", h_dwells, WIDTH, HEIGHT);

    free(h_dwells);

    printf("\nAverage GPU time over %d runs (excluding 1st run): %.6f ms\n", RUNS - 1, total_time / (RUNS - 1));
    printf("Global peak GPU memory usage over all runs: %.2f MB\n", global_peak_mem / (1024.0 * 1024.0));

    cudaEventDestroy(start);
    cudaEventDestroy(stop);

    return 0;
}